from langchain_core.tools import tool
import json
from langchain_deepseek import ChatDeepSeek
from langchain_core.messages import HumanMessage

@tool
def add(a: int, b: int) -> int:
    """Add two numbers together"""
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtract two numbers"""
    return a - b


llm = ChatDeepSeek(model="deepseek-chat", temperature=0)

# 将工具绑定到LLM上
llm_with_tools = llm.bind_tools([add, subtract])

query = "8加上15等于多少？"
result = llm_with_tools.invoke(query) # invoke方法返回一个LLMResult对象

messages = [HumanMessage(content=query)]
messages.append(result)

available_tools = {"add": add, "subtract": subtract}

# 使用工具, 获取工具的返回值
for tool_call in result.tool_calls:
    selected_tool = available_tools[tool_call["name"].lower()]
    tool_msg = selected_tool.invoke(tool_call)
    messages.append(tool_msg)

for message in messages:
     # ensure_ascii=False 表示不使用ASCII编码
    print(json.dumps(message.model_dump(), indent=4, ensure_ascii=False))

new_result = llm_with_tools.invoke(messages)
print(new_result.content)
